package com.algorithm.example.tree;


import static com.algorithm.example.tree.RedBlackTree.Color.BLACK;
import static com.algorithm.example.tree.RedBlackTree.Color.RED;


/**
 * <p>
 * <h3>红黑树</h3>
 * <li>1. 所有节点都有两种颜色：红:red_circle:、黑:black_circle</li>
 * <li>2. 所有 null 视为黑色:black_circle:</li>
 * <li>3. 红色:red_circle:节点不能相邻</li>
 * <li>4. 根节点是黑色:black_circle:</li>
 * <li>5. 从根到任意一个叶子节点，路径中的黑色:black_circle:节点数一样</li>
 * </p>
 *
 * @author lzt
 * @version 1.0
 * @since 2024/2/16 16:39
 */
public class RedBlackTree {
    enum Color {
        RED, BLACK;
    }

    RedBlackNode root;

    static class RedBlackNode {
        int key;
        Object value;
        RedBlackNode left;
        RedBlackNode right;
        RedBlackNode parent;
        Color color = RED;

        public RedBlackNode(int key) {
            this.key = key;
        }

        public RedBlackNode(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public RedBlackNode(int key, Object value, RedBlackNode left, RedBlackNode right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
            if (left != null) {
                left.parent = this;
            }
            if (right != null) {
                right.parent = this;
            }
        }

        @Override
        public String toString() {
            return "\nRedBlackNode{" +
                    "key=" + key +
                    ", value=" + value +
                    ", left=" + left +
                    ", right=" + right +
                    ", color=" + color +
                    '}';
        }

        /**
         * 是否是左孩子
         *
         * @return true || false
         */
        boolean isLeftChild() {
            if (parent == null) {
                return false;
            }
            return parent.left == this;
        }

        // 叔叔
        RedBlackNode uncle() {
            if (this.parent == null || this.parent.parent == null) {
                return null;
            }
            RedBlackNode grandfather = this.parent.parent;
            if (parent.isLeftChild()) {
                return grandfather.right;
            } else {
                return grandfather.left;
            }
        }


        /**
         * 兄弟
         *
         * @return 兄弟节点
         */
        RedBlackNode sibling() {
            if (parent == null) {
                return null;
            }
            if (isLeftChild()) {
                return parent.right;
            } else {
                return parent.left;
            }
        }
    }

    // 判断红
    boolean isRed(RedBlackNode node) {
        return node != null && node.color == RED;
    }

    // 判断黑
    boolean isBlack(RedBlackNode node) {
        return node == null || node.color == BLACK;
    }


    //右旋 1．parent的处理 2．旋转后新根的父子关系
    private void rightRotate(RedBlackNode pink) {
        RedBlackNode parent = pink.parent;
        RedBlackNode yellow = pink.left;
        RedBlackNode green = yellow.right;
        if (green != null) {
            green.parent = pink;
        }
        yellow.right = pink;
        yellow.parent = parent;
        pink.left = green;
        pink.parent = yellow;
        if (parent == null) {
            root = yellow;
        } else if (parent.left == pink) {
            parent.left = yellow;
        } else {
            parent.right = yellow;
        }
    }

    // 左旋
    private void leftRotate(RedBlackNode pink) {
        RedBlackNode parent = pink.parent;
        RedBlackNode yellow = pink.right;
        RedBlackNode green = yellow.left;
        if (green != null) {
            green.parent = pink;
        }
        yellow.left = pink;
        yellow.parent = parent;
        pink.right = green;
        pink.parent = yellow;
        if (parent == null) {
            root = yellow;
        } else if (parent.left == pink) {
            parent.left = yellow;
        } else {
            parent.right = yellow;
        }
    }


    /**
     * 新增或更新
     * <br>
     * 正常增、遇到红红不平衡进行调整
     *
     * @param key   - 键
     * @param value - 值
     */
    public void put(int key, Object value) {
        RedBlackNode p = root;
        RedBlackNode parent = null;
        while (p != null) {
            parent = p;
            if (p.key > key) {
                p = p.left;
            } else if (p.key < key) {
                p = p.right;
            } else {
                p.value = value;
                return;
            }
        }
        RedBlackNode insertNode = new RedBlackNode(key, value);
        if (root == null) {
            root = insertNode;
        } else if (key > parent.key) {
            parent.right = insertNode;
            insertNode.parent = parent;
        } else {
            parent.left = insertNode;
            insertNode.parent = parent;
        }
        fixRedRed(insertNode);
    }

    //
//    private void fixRedRed(RedBlackNode node) {
//        // case 1：插入节点为根节点，将根节点变黑:black_circle:
//        if (node == root) {
//            node.color = BLACK;
//            return;
//        }
//
//        // case 2：插入节点的父亲若为黑色:black_circle:，树的红黑性质不变，无需调整
//        if (isBlack(node.parent)) {
//            return;
//        }
//
//        //父亲为红色  ||  触发红红相邻
//        // case 3：叔叔为红色:red_circle
//        RedBlackNode uncle = node.uncle();
//        RedBlackNode parent = node.parent;
//        RedBlackNode grandfather = node.parent.parent;
//        while (isRed(uncle)) {
//            parent.color = BLACK;
//            uncle.color = BLACK;
//            grandfather.color = RED;
//
//            if (grandfather == root) {
//                grandfather.color = BLACK;
//                return;
//            }
//
//            // case 2：插入节点的父亲若为黑色:black_circle:，树的红黑性质不变，无需调整
//            if (isBlack(grandfather.parent)) {
//                return;
//            }
//
//            uncle = grandfather.uncle();
//            parent = grandfather.parent;
//            grandfather = grandfather.parent.parent;
//
//        }
//
//        // case 4：叔叔为黑色:black_circle:
//        if (isBlack(uncle) && isRed(parent)) {
//            if (parent.isLeftChild() && node.isLeftChild()) {
//                // LL
//                parent.color = BLACK;
//                grandfather.color = RED;
//                rightRotate(grandfather);
//            } else if (parent.isLeftChild() && !node.isLeftChild()) {
//                // LR -> 左旋变成ll
//                leftRotate(parent);
//                parent.color = BLACK;
//                grandfather.color = RED;
//                rightRotate(grandfather);
//            } else if (!parent.isLeftChild() && !node.isLeftChild()) {
//                // RR
//                parent.color = BLACK;
//                grandfather.color = RED;
//                leftRotate(grandfather);
//            } else {
//                // RL
//                rightRotate(parent);
//                parent.color = BLACK;
//                grandfather.color = RED;
//                leftRotate(grandfather);
//            }
//        }
//
//
//    }
    private void fixRedRed(RedBlackNode x) {
        while (true) {
            // case 1 插入节点是根节点，变黑即可
            if (x == root) {
                x.color = BLACK;
                return;
            }
            // case 2 插入节点父亲是黑色，无需调整
            if (isBlack(x.parent)) {
                return;
            }

            RedBlackNode parent = x.parent;
            RedBlackNode uncle = x.uncle();
            RedBlackNode grandparent = parent.parent;

            // case 3 当红红相邻，叔叔为红时
            if (isRed(uncle)) {
                parent.color = BLACK;
                uncle.color = BLACK;
                grandparent.color = RED;
                x = grandparent; // 更新当前节点为祖父节点，继续检查祖父节点是否满足红黑树性质
            } else { // case 4 当红红相邻，叔叔为黑时
                if (parent.isLeftChild() && x.isLeftChild()) { // LL
                    parent.color = BLACK;
                    grandparent.color = RED;
                    rightRotate(grandparent);
                    x = grandparent; // 更新当前节点为祖父节点，继续检查祖父节点是否满足红黑树性质
                } else if (parent.isLeftChild()) { // LR
                    leftRotate(parent);
                    x.color = BLACK;
                    grandparent.color = RED;
                    rightRotate(grandparent);
                    x = grandparent; // 更新当前节点为祖父节点，继续检查祖父节点是否满足红黑树性质
                } else if (!x.isLeftChild()) { // RR
                    parent.color = BLACK;
                    grandparent.color = RED;
                    leftRotate(grandparent);
                    x = grandparent; // 更新当前节点为祖父节点，继续检查祖父节点是否满足红黑树性质
                } else { // RL
                    rightRotate(parent);
                    x.color = BLACK;
                    grandparent.color = RED;
                    leftRotate(grandparent);
                    x = grandparent; // 更新当前节点为祖父节点，继续检查祖父节点是否满足红黑树性质
                }
            }
        }
    }

    /**
     * 删除
     *
     * @param key - 待删键
     */
    public void remove(int key) {
        RedBlackNode p = root;
        while (p != null) {
            if (p.key > key) {
                p = p.left;
            } else if (p.key < key) {
                p = p.right;
            } else {
                break;
            }
        }

        if (p == null) {
            return;
        }

        doRemove(p);

    }

    private void doRemove(RedBlackNode delete) {
        RedBlackNode replaced = findReplaced(delete);
        RedBlackNode parent = delete.parent;

        // 没有孩子
        if (replaced == null) {
            // case 1 删除的根节点
            if (delete == root) {
                root = null;
            } else {
                if (isBlack(delete)) {
                    // 双黑调整
                    fixDoubleBlack(delete);
                } else {
                    // 红色叶子，无需处理
                }
                if (delete.isLeftChild()) {
                    parent.left = null;
                } else {
                    parent.right = null;
                }
                delete.parent = null;
            }
            return;
        }
        // 有一个孩子
        if (delete.left == null || delete.right == null) {
            // case 1 删除的根节点
            if (delete == root) {
                root.key = replaced.key;
                root.value = replaced.value;
                root.left = root.right = null;
            } else {
                if (delete.isLeftChild()) {
                    parent.left = replaced;
                } else {
                    parent.right = replaced;
                }
                replaced.parent = parent;
                delete.left = delete.right = delete.parent = null;
                if (isBlack(delete) && isBlack(replaced)) {
                    // 复杂处理 @TODO 实际不会有这种情况 因为只有一个孩子时 被删除节点是黑色 那么剩余节点只能是红色不会触发双黑
                    fixDoubleBlack(replaced);
                } else {
                    // case 2 删除是黑 剩下是红
                    replaced.color = BLACK;
                }
            }
            return;
        }

        // case 0 有两个孩子 => 有一个孩子 或 没有孩子
        int t = delete.key;
        delete.key = replaced.key;
        replaced.key = t;

        Object v = delete.value;
        delete.value = replaced.value;
        replaced.value = v;
        doRemove(replaced);
    }

    /**
     * 双黑调整
     *
     * @param x - 被调整的节点
     */
    private void fixDoubleBlack(RedBlackNode x) {
        if (x == root) {
            return;
        }

        RedBlackNode parent = x.parent;
        RedBlackNode sibling = x.sibling();

        // case 3 调整节点的兄弟为红:red_circle:，此时两个侄子定为黑:black_circle:
        if (isRed(sibling)) {
            if (x.isLeftChild()) {
                leftRotate(parent);
            } else {
                rightRotate(parent);
            }
            parent.color = RED;
            sibling.color = BLACK;
            fixDoubleBlack(x);
            return;
        }
        if (sibling != null) {
            // case 4：被调整节点的兄弟为黑:black_circle:，两个侄子都为黑:black_circle:
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                sibling.color = RED;
                if (isRed(parent)) {
                    parent.color = BLACK;
                } else {
                    fixDoubleBlack(parent);
                }
            }
            // case 5 调整节点的兄弟为黑:black_circle:，至少一个红:red_circle:侄子
            else {
                //LL
                if (sibling.isLeftChild() && isRed(sibling.left)) {
                    rightRotate(parent);
                    sibling.left.color = BLACK;
                    sibling.color = parent.color;
                } else

                    //lR
                    if (sibling.isLeftChild() && isRed(sibling.right)) {
                        sibling.right.color = parent.color;
                        leftRotate(sibling);
                        rightRotate(parent);
                    } else

                        // RR
                        if (!sibling.isLeftChild() && isRed(sibling.right)) {
                            parent.color = BLACK;
                            leftRotate(parent);
                            sibling.right.color = parent.color;
                        }

                        // RL
                        else {
                            sibling.left.color = parent.color;
                            leftRotate(sibling);
                            rightRotate(parent);
                        }
                parent.color = BLACK;

            }

        } else {
            // @TODO 实际也不会出现，触发双黑后，兄弟节点不会为 null
            fixDoubleBlack(parent);
        }
    }

    /**
     * 查找剩余节点
     *
     * @param delete - 被删除的节点
     * @return - 返回剩余节点
     */
    private RedBlackNode findReplaced(RedBlackNode delete) {
        if (delete.left == null && delete.right == null) {
            return null;
        }

        if (delete.left == null) {
            return delete.right;
        }

        if (delete.right == null) {
            return delete.left;
        }

        // 寻找最小值
        RedBlackNode node = delete.right;
        while (node.left != null) {
            node = node.left;
        }
        return node;
    }


    @Override
    public String toString() {
        return "RedBlackTree{" +
                "root=" + root +
                '}';
    }
}
